from llama_index.core import SimpleDirectoryReader
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

from llama_index.core import SimpleDirectoryReader
from llama_index.core.node_parser import SimpleNodeParser
from llama_index.core import  GPTVectorStoreIndex,VectorStoreIndex
from llama_index.llms import openai_like
from llama_index.core import Settings
from llama_index.llms.ollama import Ollama
from llama_index.embeddings.huggingface import HuggingFaceEmbedding  # HuggingFaceEmbedding:用于将文本转换为词向量
from llama_index.llms.huggingface import HuggingFaceLLM  # HuggingFaceLLM：用于运行Hugging Face的预训练语言模型
from llama_index.core import Settings,SimpleDirectoryReader,VectorStoreIndex
import chromadb
from llama_index.embeddings.dashscope import DashScopeEmbedding
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import StorageContext, load_index_from_storage
from llama_index.llms.deepseek  import DeepSeek
from llama_index.embeddings.fastembed import FastEmbedEmbedding

from llama_index.core import QueryBundle

# import NodeWithScore
from llama_index.core.schema import NodeWithScore

# Retrievers
from llama_index.core.retrievers import (
    BaseRetriever,
    VectorIndexRetriever,
    KeywordTableSimpleRetriever,
)
    # 连接Chroma数据库


llm = DeepSeek(model="deepseek-chat", api_key="sk-605e60a1301040759a821b6b677556fb")
Settings.llm = llm
 
from zhipuai import ZhipuAI
from llama_index.embeddings.zhipuai import ZhipuAIEmbedding

embeddings = ZhipuAIEmbedding(
    model="embedding-2",
    api_key="f387f5e4837d4e4bba6d267682a957c9.PmPiTw8qVlsI2Oi5"
    # With the `embedding-3` class
    # of models, you can specify the size
    # of the embeddings you want returned.
    # dimensions=1024
)
Settings.embed_model=embeddings

from typing import Optional

from llama_index.core.tools import FunctionTool
from pydantic import BaseModel

# we will store booking under random IDs
bookings = {}


# we will represent and track the state of a booking as a Pydantic model
class Booking(BaseModel):
    name: Optional[str] = None
    email: Optional[str] = None
    phone: Optional[str] = None
    date: Optional[str] = None
    time: Optional[str] = None


def get_booking_state(user_id: str) -> str:
    """Get the current state of a booking for a given booking ID."""
    try:
        return str(bookings[user_id].dict())
    except:
        return f"Booking ID {user_id} not found"


def update_booking(user_id: str, property: str, value: str) -> str:
    """Update a property of a booking for a given booking ID. Only enter details that are explicitly provided."""
    booking = bookings[user_id]
    setattr(booking, property, value)
    return f"Booking ID {user_id} updated with {property} = {value}"


def create_booking(user_id: str) -> str:
    """Create a new booking and return the booking ID."""
    bookings[user_id] = Booking()
    return "Booking created, but not yet confirmed. Please provide your name, email, phone, date, and time."


def confirm_booking(user_id: str) -> str:
    """Confirm a booking for a given booking ID."""
    booking = bookings[user_id]

    if booking.name is None:
        raise ValueError("Please provide your name.")

    if booking.email is None:
        raise ValueError("Please provide your email.")

    if booking.phone is None:
        raise ValueError("Please provide your phone number.")

    if booking.date is None:
        raise ValueError("Please provide the date of your booking.")

    if booking.time is None:
        raise ValueError("Please provide the time of your booking.")

    return f"Booking ID {user_id} confirmed!"


# create tools for each function
get_booking_state_tool = FunctionTool.from_defaults(fn=get_booking_state)
update_booking_tool = FunctionTool.from_defaults(fn=update_booking)
create_booking_tool = FunctionTool.from_defaults(
    fn=create_booking, return_direct=True
)
confirm_booking_tool = FunctionTool.from_defaults(
    fn=confirm_booking, return_direct=True
)


from llama_index.core.llms import ChatMessage
from llama_index.core.agent.workflow import FunctionAgent
from llama_index.core.workflow import Context

 

user = "user123"
system_prompt = f"""You are now connected to the booking system and helping {user} with making a booking.
Only enter details that the user has explicitly provided.
Do not make up any details.
"""

agent = FunctionAgent(
    tools=[
        get_booking_state_tool,
        update_booking_tool,
        create_booking_tool,
        confirm_booking_tool,
    ],
    llm=llm,
    system_prompt=system_prompt,
)

# create a context for the agent to hold the state/history of a session
ctx = Context(agent)

from llama_index.core.agent.workflow import AgentStream, ToolCallResult
 

async def Run():
    handler = agent.run(
    "Hello! I would like to make a booking, around 5pm?", ctx=ctx
)

    async for ev in handler.stream_events():
        if isinstance(ev, AgentStream):
            print(f"{ev.delta}", end="", flush=True)
        elif isinstance(ev, ToolCallResult):
            print(
                f"\nCall {ev.tool_name} with {ev.tool_kwargs}\nReturned: {ev.tool_output}"
            )

    response = await handler
    print(str(response))
    handler = agent.run(
    "Sure! My name is Logan, and my email is test@gmail.com?", ctx=ctx
)

    async for ev in handler.stream_events():
        if isinstance(ev, AgentStream):
            print(f"{ev.delta}", end="", flush=True)
        elif isinstance(ev, ToolCallResult):
            print(
                f"\nCall {ev.tool_name} with {ev.tool_kwargs}\nReturned: {ev.tool_output}"
            )

    response = await handler
    print(str(response))
    handler = agent.run(
    "Right! My phone number is 1234567890, the date of the booking is April 5, at 5pm.",
    ctx=ctx,
)

    async for ev in handler.stream_events():
        if isinstance(ev, AgentStream):
            print(f"{ev.delta}", end="", flush=True)
        elif isinstance(ev, ToolCallResult):
            print(
                f"\nCall {ev.tool_name} with {ev.tool_kwargs}\nReturned: {ev.tool_output}"
            )

    response = await handler

    print(str(response))
    print(bookings["user123"])



import asyncio
print(__name__)
if __name__=="__main__":
    asyncio.run(Run())